#include <unittest/unittest.h>

#ifndef _MSC_VER
#define CUSP_USE_TEXTURE_MEMORY
#endif

#include <cusp/multiply.h>

#include <cusp/linear_operator.h>
#include <cusp/print.h>
#include <cusp/gallery/poisson.h>
#include <cusp/gallery/random.h>

#include <cusp/array2d.h>
#include <cusp/coo_matrix.h>
#include <cusp/csr_matrix.h>
#include <cusp/dia_matrix.h>
#include <cusp/ell_matrix.h>
#include <cusp/hyb_matrix.h>
#include <cusp/permutation_matrix.h>

/////////////////////////////////////////
// Sparse Matrix-Matrix Multiplication //
/////////////////////////////////////////

template <typename SparseMatrixType, typename DenseMatrixType>
void CompareSparseMatrixMatrixMultiply(DenseMatrixType A, DenseMatrixType B)
{
    DenseMatrixType C;
    cusp::multiply(A, B, C);

    SparseMatrixType _A(A), _B(B), _C;
    cusp::multiply(_A, _B, _C);
    
    ASSERT_EQUAL(C == DenseMatrixType(_C), true);
    
    typename SparseMatrixType::view _Aview(_A), _Bview(_B), _Cview(_C);
    cusp::multiply(_Aview, _Bview, _Cview);
    
    ASSERT_EQUAL(C == DenseMatrixType(_Cview), true);
}

template <typename TestMatrix>
void TestSparseMatrixMatrixMultiply(void)
{
    cusp::array2d<float,cusp::host_memory> A(3,2);
    A(0,0) = 1.0; A(0,1) = 2.0;
    A(1,0) = 3.0; A(1,1) = 0.0;
    A(2,0) = 5.0; A(2,1) = 6.0;
    
    cusp::array2d<float,cusp::host_memory> B(2,4);
    B(0,0) = 0.0; B(0,1) = 2.0; B(0,2) = 3.0; B(0,3) = 4.0;
    B(1,0) = 5.0; B(1,1) = 0.0; B(1,2) = 0.0; B(1,3) = 8.0;

    cusp::array2d<float,cusp::host_memory> C(2,2);
    C(0,0) = 0.0; C(0,1) = 0.0;
    C(1,0) = 3.0; C(1,1) = 5.0;
    
    cusp::array2d<float,cusp::host_memory> D(2,1);
    D(0,0) = 2.0;
    D(1,0) = 3.0;
    
    cusp::array2d<float,cusp::host_memory> E(2,2);
    E(0,0) = 0.0; E(0,1) = 0.0;
    E(1,0) = 0.0; E(1,1) = 0.0;
    
    cusp::array2d<float,cusp::host_memory> F(2,3);
    F(0,0) = 0.0; F(0,1) = 1.5; F(0,2) = 3.0;
    F(1,0) = 0.5; F(1,1) = 0.0; F(1,2) = 0.0;
    
    cusp::array2d<float,cusp::host_memory> G;
    cusp::gallery::poisson5pt(G, 4, 6);

    cusp::array2d<float,cusp::host_memory> H;
    cusp::gallery::poisson5pt(H, 8, 3);

    cusp::array2d<float,cusp::host_memory> I;
    cusp::gallery::random(24, 24, 150, I);
    
    cusp::array2d<float,cusp::host_memory> J;
    cusp::gallery::random(24, 24, 50, J);

    cusp::array2d<float,cusp::host_memory> K;
    cusp::gallery::random(24, 12, 20, K);
 
    thrust::host_vector< cusp::array2d<float,cusp::host_memory> > matrices;
    matrices.push_back(A);
    matrices.push_back(B);
    matrices.push_back(C);
    matrices.push_back(D);
    matrices.push_back(E);
    matrices.push_back(F);
    matrices.push_back(G);
    matrices.push_back(H);
    matrices.push_back(I);
    matrices.push_back(J);
    matrices.push_back(K);

    // test matrix multiply for every pair of compatible matrices
    for(size_t i = 0; i < matrices.size(); i++)
    {
        const cusp::array2d<float,cusp::host_memory>& left = matrices[i];
        for(size_t j = 0; j < matrices.size(); j++)
        {
            const cusp::array2d<float,cusp::host_memory>& right = matrices[j];

            if (left.num_cols == right.num_rows)
                CompareSparseMatrixMatrixMultiply<TestMatrix>(left, right);
        }
    }

}
DECLARE_SPARSE_MATRIX_UNITTEST(TestSparseMatrixMatrixMultiply);


/////////////////////////////////////////
// Sparse Matrix-Vector Multiplication //
/////////////////////////////////////////

template <typename SparseMatrixType, typename DenseMatrixType>
void CompareSparseMatrixVectorMultiply(DenseMatrixType A)
{
    typedef typename SparseMatrixType::memory_space MemorySpace;

    // setup reference input
    cusp::array1d<float, cusp::host_memory> x(A.num_cols);
    cusp::array1d<float, cusp::host_memory> y(A.num_rows, 10);
    for(size_t i = 0; i < x.size(); i++)
        x[i] = i % 10;

    // compute reference output
    cusp::multiply(A, x, y);
  
    // test container
    {
      SparseMatrixType _A(A);
      cusp::array1d<float, MemorySpace> _x(x);
      cusp::array1d<float, MemorySpace> _y(A.num_rows, 10);

      cusp::multiply(_A, _x, _y);
      
      ASSERT_EQUAL(_y, y);
    }

    // test matrix view
    {
      SparseMatrixType _A(A);
      cusp::array1d<float, MemorySpace> _x(x);
      cusp::array1d<float, MemorySpace> _y(A.num_rows, 10);

      typename SparseMatrixType::view _V(_A);
      cusp::multiply(_V, _x, _y);
      
      ASSERT_EQUAL(_y, y);
    }
    
    // test array view
    {
      SparseMatrixType _A(A);
      cusp::array1d<float, MemorySpace> _x(x);
      cusp::array1d<float, MemorySpace> _y(A.num_rows, 10);

      typename cusp::array1d<float, MemorySpace> _Vx(_x), _Vy(_y);
      cusp::multiply(_A, _Vx, _Vy);
      
      ASSERT_EQUAL(_Vy, y);
    }
}


// TODO use COO reference format and test larger problem sizes
template <class TestMatrix>
void TestSparseMatrixVectorMultiply()
{
    typedef typename TestMatrix::memory_space MemorySpace;

    cusp::array2d<float, cusp::host_memory> A(4,4);
    A(0,0) = 13; A(0,1) = 80; A(0,2) =  0; A(0,3) =  0; 
    A(1,0) =  0; A(1,1) = 27; A(1,2) =  0; A(1,3) =  0;
    A(2,0) = 55; A(2,1) =  0; A(2,2) = 24; A(2,3) = 42;
    A(3,0) =  0; A(3,1) = 69; A(3,2) =  0; A(3,3) = 83;
    
    cusp::array2d<float,cusp::host_memory> B(2,4);
    B(0,0) = 0.0; B(0,1) = 2.0; B(0,2) = 3.0; B(0,3) = 4.0;
    B(1,0) = 5.0; B(1,1) = 0.0; B(1,2) = 0.0; B(1,3) = 8.0;

    cusp::array2d<float,cusp::host_memory> C(2,2);
    C(0,0) = 0.0; C(0,1) = 0.0;
    C(1,0) = 3.0; C(1,1) = 5.0;
    
    cusp::array2d<float,cusp::host_memory> D(2,1);
    D(0,0) = 2.0;
    D(1,0) = 3.0;
    
    cusp::array2d<float,cusp::host_memory> E(2,2);
    E(0,0) = 0.0; E(0,1) = 0.0;
    E(1,0) = 0.0; E(1,1) = 0.0;
    
    cusp::array2d<float,cusp::host_memory> F(2,3);
    F(0,0) = 0.0; F(0,1) = 1.5; F(0,2) = 3.0;
    F(1,0) = 0.5; F(1,1) = 0.0; F(1,2) = 0.0;
    
    cusp::array2d<float,cusp::host_memory> G;
    cusp::gallery::poisson5pt(G, 4, 6);

    cusp::array2d<float,cusp::host_memory> H;
    cusp::gallery::poisson5pt(H, 8, 3);

    CompareSparseMatrixVectorMultiply<TestMatrix>(A);
    CompareSparseMatrixVectorMultiply<TestMatrix>(B);
    CompareSparseMatrixVectorMultiply<TestMatrix>(C);
    CompareSparseMatrixVectorMultiply<TestMatrix>(D);
    CompareSparseMatrixVectorMultiply<TestMatrix>(E);
    CompareSparseMatrixVectorMultiply<TestMatrix>(F);
    CompareSparseMatrixVectorMultiply<TestMatrix>(G);
    CompareSparseMatrixVectorMultiply<TestMatrix>(H);
}
DECLARE_SPARSE_MATRIX_UNITTEST(TestSparseMatrixVectorMultiply);


//////////////////////////////
// General Linear Operators //
//////////////////////////////

template <class MemorySpace>
void TestMultiplyIdentityOperator(void)
{
    cusp::array1d<float, MemorySpace> x(4);
    cusp::array1d<float, MemorySpace> y(4);

    x[0] =  7.0f;   y[0] =  0.0f; 
    x[1] =  5.0f;   y[1] = -2.0f;
    x[2] =  4.0f;   y[2] =  0.0f;
    x[3] = -3.0f;   y[3] =  5.0f;

    cusp::identity_operator<float, MemorySpace> A(4,4);
    
    cusp::multiply(A, x, y);

    ASSERT_EQUAL(y[0],  7.0f);
    ASSERT_EQUAL(y[1],  5.0f);
    ASSERT_EQUAL(y[2],  4.0f);
    ASSERT_EQUAL(y[3], -3.0f);
}
DECLARE_HOST_DEVICE_UNITTEST(TestMultiplyIdentityOperator);

//////////////////////////////
// Permutation Operators    //
//////////////////////////////
template <typename SparseMatrixType, typename DenseMatrixType1, typename DenseMatrixType2>
void CompareRowPermutationMatrixMatrixMultiply(DenseMatrixType1 A, DenseMatrixType2 B)
{
    typedef typename SparseMatrixType::index_type   IndexType;
    typedef typename SparseMatrixType::memory_space MemorySpace;
    typedef cusp::permutation_matrix<IndexType, MemorySpace> PermutationMatrixType;

    DenseMatrixType2 C;
    cusp::multiply(A, B, C);

    PermutationMatrixType _A(A);
    SparseMatrixType _B(B), _C;
    cusp::multiply(_A, _B, _C);

    ASSERT_EQUAL(C == DenseMatrixType2(_C), true);
    
    typename PermutationMatrixType::view _Aview(_A);
    typename SparseMatrixType::view _Bview(_B), _Cview(_C);
    cusp::multiply(_Aview, _Bview, _Cview);
    
    ASSERT_EQUAL(C == DenseMatrixType2(_Cview), true);
}

template <typename SparseMatrixType, typename DenseMatrixType1, typename DenseMatrixType2>
void CompareColPermutationMatrixMatrixMultiply(DenseMatrixType1 A, DenseMatrixType2 B)
{
    typedef typename SparseMatrixType::index_type   IndexType;
    typedef typename SparseMatrixType::memory_space MemorySpace;
    typedef cusp::permutation_matrix<IndexType, MemorySpace> PermutationMatrixType;

    DenseMatrixType1 C;
    cusp::multiply(A, B, C);

    PermutationMatrixType _B(B);
    SparseMatrixType _A(A), _C;
    cusp::multiply(_A, _B, _C);

    ASSERT_EQUAL(C == DenseMatrixType1(_C), true);
    
    typename PermutationMatrixType::view _Bview(_B);
    typename SparseMatrixType::view _Aview(_A), _Cview(_C);
    cusp::multiply(_Aview, _Bview, _Cview);
   
    ASSERT_EQUAL(C == DenseMatrixType1(_Cview), true);
}

template <class MemorySpace>
void TestPermutationMatrixVectorMultiply(void)
{
    cusp::array1d<float, MemorySpace> x(4);
    cusp::array1d<float, MemorySpace> y(4);

    x[0] =  7.0f;   y[0] =  0.0f; 
    x[1] =  5.0f;   y[1] = -2.0f;
    x[2] =  4.0f;   y[2] =  0.0f;
    x[3] = -3.0f;   y[3] =  5.0f;

    cusp::permutation_matrix<int, MemorySpace> A(4);

    A.values[0] = 3; 
    A.values[1] = 2; 
    A.values[2] = 1; 
    A.values[3] = 0;
    
    cusp::multiply(A, x, y);

    ASSERT_EQUAL(y[0], -3.0f);
    ASSERT_EQUAL(y[1],  4.0f);
    ASSERT_EQUAL(y[2],  5.0f);
    ASSERT_EQUAL(y[3],  7.0f);
}
DECLARE_HOST_DEVICE_UNITTEST(TestPermutationMatrixVectorMultiply);

template <class MemorySpace>
void TestPermutationMatrixMatrixMultiply(void)
{
    typedef cusp::csr_matrix<int, int, MemorySpace> TestMatrix;

    cusp::array2d<int, cusp::host_memory> A(4,4);
    A(0,0) = 0; A(0,1) = 0; A(0,2) = 0; A(0,3) = 1; 
    A(1,0) = 0; A(1,1) = 0; A(1,2) = 1; A(1,3) = 0;
    A(2,0) = 0; A(2,1) = 1; A(2,2) = 0; A(2,3) = 0;
    A(3,0) = 1; A(3,1) = 0; A(3,2) = 0; A(3,3) = 0;

    cusp::array2d<int, cusp::host_memory> B(4,4);
    B(0,0) = 13; B(0,1) = 80; B(0,2) =  0; B(0,3) =  0; 
    B(1,0) =  0; B(1,1) = 27; B(1,2) =  0; B(1,3) =  0;
    B(2,0) = 55; B(2,1) =  0; B(2,2) = 24; B(2,3) = 42;
    B(3,0) =  0; B(3,1) = 69; B(3,2) =  0; B(3,3) = 83;
    
    cusp::array2d<int, cusp::host_memory> C;
    cusp::gallery::poisson5pt(C, 2, 2);

    thrust::host_vector< cusp::array2d<int, cusp::host_memory> > matrices;
    matrices.push_back(B);
    matrices.push_back(C);

    // test row permutations 
    for(size_t i = 0; i < matrices.size(); i++)
    {
        const cusp::array2d<int, cusp::host_memory>& M = matrices[i];
        CompareRowPermutationMatrixMatrixMultiply<TestMatrix>(A, M);
    }

    // test column permutations 
    for(size_t i = 0; i < matrices.size(); i++)
    {
        const cusp::array2d<int, cusp::host_memory>& M = matrices[i];
        CompareColPermutationMatrixMatrixMultiply<TestMatrix>(M, A);
    }

}
DECLARE_HOST_DEVICE_UNITTEST(TestPermutationMatrixMatrixMultiply);

void TestDevicePermutationMatrixMatrixMultiply(void)
{
    cusp::coo_matrix<int,float,cusp::device_memory> A;
    cusp::gallery::poisson5pt(A, 4, 4);

    cusp::permutation_matrix<int,cusp::device_memory> P(A.num_rows);
    P.random_permutation();

    {
      cusp::coo_matrix<int,float,cusp::device_memory> B;
      cusp::multiply(A,P,B);

      ASSERT_EQUAL(B.is_sorted_by_row_and_column(), true);
    }

    {
      cusp::coo_matrix<int,float,cusp::device_memory> B;
      cusp::multiply(P,A,B);

      ASSERT_EQUAL(B.is_sorted_by_row_and_column(), true);
    }
}
DECLARE_UNITTEST(TestDevicePermutationMatrixMatrixMultiply);

